package cn.bugstack.config;

import cn.bugstack.http.OllamaController;
import org.springframework.ai.ollama.OllamaChatClient;
import org.springframework.ai.ollama.OllamaEmbeddingClient;
import org.springframework.ai.ollama.api.OllamaApi;
import org.springframework.ai.ollama.api.OllamaModel;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.PgVectorStore;
import org.springframework.ai.vectorstore.SimpleVectorStore;
import org.springframework.beans.factory.annotation.Configurable;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.jdbc.core.JdbcTemplate;

/**
 * @Classname OllamaConfig
 * @Description TODO
 * @Date 2025/10/16 11:25
 * @Author 勾新杰
 */
@Configuration
public class OllamaConfig {

    @Bean
    public OllamaApi ollamaApi(@Value("${spring.ai.ollama.base-url}") String baseUrl) {
        return new OllamaApi(baseUrl);
    }

    @Bean
    public OllamaChatClient ollamaChatClient(OllamaApi ollamaApi) {
        return new OllamaChatClient(ollamaApi);
    }

    @Bean
    public TokenTextSplitter tokenTextSplitter() {
        return new TokenTextSplitter();
    }

    @Bean
    public OllamaEmbeddingClient ollamaEmbeddingClient(OllamaApi ollamaApi, @Value("${spring.ai.ollama.embedding.model}") String embeddingModel) {
        return new OllamaEmbeddingClient(ollamaApi).withDefaultOptions(
                OllamaOptions.create().withModel(embeddingModel));
    }

    @Bean
    public SimpleVectorStore simpleVectorStore(OllamaEmbeddingClient ollamaEmbeddingClient) {
        return new SimpleVectorStore(ollamaEmbeddingClient);
    }

    @Bean
    public PgVectorStore pgVectorStore(OllamaEmbeddingClient ollamaEmbeddingClient, JdbcTemplate jdbcTemplate){
        return new PgVectorStore(jdbcTemplate, ollamaEmbeddingClient);
    }
}
